#!/usr/bin/env python

# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Train a Fast R-CNN network on a region of interest database."""


import _init_paths
from fast_rcnn.train import get_training_roidb, train_net
from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list, get_output_dir
from datasets.factory import get_imdb
import datasets.imdb
import argparse
import pprint
import numpy as np
import sys
import os
import caffe

# usage:
# /tools/train_net.py --gpu ${GPU_ID} \
#   --solver models/${PT_DIR}/${NET}/faster_rcnn_end2end/solver.prototxt \
#   --weights data/imagenet_models/${NET}.v2.caffemodel \
#   --imdb ${TRAIN_IMDB} \
#   --iters ${ITERS} \
#   --cfg experiments/cfgs/faster_rcnn_end2end.yml \
#   ${EXTRA_ARGS}

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
    parser.add_argument('--gpu', dest='gpu_id',
                        help='GPU device id to use [0]',
                        default=0, type=int)
    parser.add_argument('--solver', dest='solver',
                        help='solver prototxt',
                        default=None, type=str)
    parser.add_argument('--iters', dest='max_iters',
                        help='number of iterations to train',
                        default=40000, type=int)
    parser.add_argument('--weights', dest='pretrained_model',
                        help='initialize with pretrained model weights',
                        default=None, type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file',
                        default=None, type=str)
    parser.add_argument('--imdb', dest='imdb_name',
                        help='dataset to train on',
                        default='voc_2007_trainval', type=str)
    parser.add_argument('--rand', dest='randomize',
                        help='randomize (do not use a fixed seed)',
                        action='store_true')
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    return args

def combined_roidb(imdb_names, devkit_path, fold_name):
    def get_roidb(imdb_name):
        imdb = get_imdb(imdb_name, devkit_path, fold_name)
        print 'Loaded dataset `{:s}` for training'.format(imdb.name)
        imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
        print 'Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD)
        roidb = get_training_roidb(imdb)
        return roidb

    roidbs = [get_roidb(s) for s in imdb_names.split('+')]
    roidb = roidbs[0]
    if len(roidbs) > 1:
        for r in roidbs[1:]:
            roidb.extend(r)
        imdb = datasets.imdb.imdb(imdb_names)
    else:
        imdb = get_imdb(imdb_names, devkit_path, fold_name)
    return imdb, roidb



def train_process(gpu_id,
                  devkit_path,
                  fold_name,
                  output_path,
                  pretrained_model,
                  train_imdb,
                  solver,
                  max_iters
                  ):

    cfg_from_file('experiments/cfgs/faster_rcnn_end2end.yml')
    cfg.GPU_ID = gpu_id

    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)

    caffe.set_mode_gpu()
    caffe.set_device(gpu_id)

    imdb, roidb = combined_roidb(train_imdb, devkit_path=devkit_path, fold_name=fold_name)

    print '{:d} roidb entries'.format(len(roidb))
    print 'Output will be saved to `{:s}`'.format(output_path)
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    train_net(solver, roidb, output_path,
              pretrained_model=pretrained_model,
              max_iters=max_iters)

if __name__ == '__main__':
    print('should not call main function')
    # args = parse_args()
    #
    # print('Called with args:')
    # print type(args)
    # print(args)
    #
    # if args.cfg_file is not None:
    #     cfg_from_file(args.cfg_file)
    # if args.set_cfgs is not None:
    #     set_cfgs_not_none
    #     cfg_from_list(args.set_cfgs)
    #
    # cfg.GPU_ID = args.gpu_id
    #
    # print('Using config:')
    # pprint.pprint(cfg)
    #
    # if not args.randomize:
    #     # fix the random seeds (numpy and caffe) for reproducibility
    #     np.random.seed(cfg.RNG_SEED)
    #     caffe.set_random_seed(cfg.RNG_SEED)
    #
    # # set up caffe
    # caffe.set_mode_gpu()
    # # caffe.set_mode_cpu()
    # caffe.set_device(args.gpu_id)
    #
    # imdb, roidb = combined_roidb(args.imdb_name)
    # print '{:d} roidb entries'.format(len(roidb))
    #
    # with open('/ssd1/ori.txt', 'w') as f:
    #     for roi in roidb:
    #         f.write(roi['image'] + '\n')
    #
    # # dj_debug
    #
    # output_dir = get_output_dir(imdb)
    #
    # print 'Output will be saved to `{:s}`'.format(output_dir)
    #
    #
    # train_net(args.solver, roidb, output_dir,
    #           pretrained_model=args.pretrained_model,
    #           max_iters=args.max_iters)
